import os
import zipfile

import cv2 as cv
from tqdm import tqdm

from mtcnn.detector import Detector
from utils import ensure_folder


def extract(filename):
    print('Extracting {}...'.format(filename))
    zip_ref = zipfile.ZipFile(filename, 'r')
    zip_ref.extractall('data')
    zip_ref.close()


if __name__ == "__main__":
    IMG_FOLDER = 'data/WIDER_val/images'
    RET_FOLDER = 'widerface_evaluate/preds'

    detector = Detector('mnet', 'cuda')

    ensure_folder(RET_FOLDER)

    dirs = [d for d in os.listdir(IMG_FOLDER)]
    dirs = sorted(dirs)

    for d in tqdm(dirs):
        src_dir = os.path.join(IMG_FOLDER, d)
        dst_dir = os.path.join(RET_FOLDER, d)
        ensure_folder(dst_dir)

        files = [f for f in os.listdir(src_dir) if f.lower().endswith('.jpg')]
        files = sorted(files)

        for f in files:
            image_name = f.replace('.jpg', '')
            lines = ['{}\n'.format(image_name)]
            src_file = os.path.join(src_dir, f)
            dst_file = os.path.join(dst_dir, image_name + '.txt')

            img = cv.imread(src_file)
            try:
                # dets, landms = detector.detect_faces(img, confidence_threshold=0.02, top_k=5000, nms_threshold=0.4,
                #                                      keep_top_k=750)
                dets, landms = detector.detect_faces(img, min_face_size=10.0, thresholds=[0.02, 0.02, 0.02],
                                                     nms_thresholds=[0.4, 0.4, 0.4])
            except ValueError as err:
                dets = []
                print(err)
            lines.append('{}\n'.format(len(dets)))

            for face in dets:
                left = int(face[0])
                top = int(face[1])
                width = int(face[2]) - left
                height = int(face[3]) - top
                score = face[4]
                lines.append('{} {} {} {} {:.3f}\n'.format(left, top, width, height, score))
            lines.append('\n')

            with open(dst_file, 'w') as fp:
                fp.writelines(lines)

            # print(dst_file)
